{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch as t\n",
    "from torch import nn, optim\n",
    "import numpy as np\n",
    "from model import *\n",
    "from multiprocess import Process, Queue\n",
    "import argparse\n",
    "from tqdm import tqdm\n",
    "from torch.nn import BCEWithLogitsLoss\n",
    "from dgl import NID, EID\n",
    "from dgl.dataloading import GraphDataLoader\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from utils import parse_arguments\n",
    "from utils import *#load_ogb_dataset,load_hmdd_dataset, evaluate_hits,standardize_dir\n",
    "from sampler import SEALData,SEALData_case_study\n",
    "from model import AE, GCN, DGCNN\n",
    "from logger import LightLogging\n",
    "import pdb\n",
    "import os.path as osp\n",
    "from eval_irgat import *\n",
    "\n",
    "np.random.seed(1337)\n",
    "device = t.device('cuda:0')#('cuda' if t.cuda.is_available() else 'cpu')\n",
    "\n",
    "def save2File(input, out_path):\n",
    "    if str(type(input)).find('ndarray') < 0:\n",
    "        input = np.array(input)\n",
    "    df = pd.DataFrame(input)\n",
    "    df.to_csv(out_path, header=False, index=False)\n",
    "    \n",
    "def adjMatrix2list(csv_path, dis_id):\n",
    "    vals = pd.read_csv(csv_path,header = None).values\n",
    "    n_miRNA = vals.shape[0]\n",
    "    n_disease = vals.shape[1]\n",
    "\n",
    "    # get possivite known association, put into pos_assoc\n",
    "    tmp_res1 = np.where(vals == 1)\n",
    "    pos_assoc = list(zip(tmp_res1[0], tmp_res1[1]))\n",
    "\n",
    "    \n",
    "    test_lbl = vals[:,dis_id].tolist()\n",
    "    test_assoc = list(zip(range(0,n_miRNA), [dis_id]*n_miRNA))\n",
    "    \n",
    "    return pos_assoc, test_assoc,test_lbl, n_miRNA, n_disease\n",
    "\n",
    "def gen_train_test(dis_id, args):\n",
    "    '''train(all known associations)\n",
    "    '''\n",
    "    adj_path = args.data_dir + 'm-d.csv'\n",
    "    train_pos_assoc, test_assoc,test_lbl, n_miRNA, n_disease = adjMatrix2list(adj_path, dis_id)\n",
    "    print('len(train_pos_assoc:', len(train_pos_assoc), 'len(test_assoc):', len(test_assoc))\n",
    "         \n",
    "    train_file_path    =  args.fold_dir + 'train.csv'\n",
    "    test_file_path     = args.fold_dir + 'test.csv'\n",
    "    test_lbl_file_path = args.fold_dir + 'test_lbl.csv'\n",
    "    \n",
    "    #if not osp.exists(train_file_path): \n",
    "    save2File(train_pos_assoc, train_file_path)\n",
    "    \n",
    "    #if not osp.exists(test_file_path):        \n",
    "    save2File(test_assoc, test_file_path)\n",
    "              \n",
    "    #if not osp.exists(test_lbl_file_path):        \n",
    "    save2File(test_lbl, test_lbl_file_path)\n",
    "    \n",
    "    return\n",
    "def case_study(dis_id,args):\n",
    "    gen_train_test(dis_id,args)\n",
    "    prefix=args.fold_dir\n",
    "    graph, split_edge = load_hmdd_dataset_case_study(args.data_dir,prefix,0.)\n",
    "    if args.sim_type == 'none':\n",
    "        features = None\n",
    "        m_d_path = args.data_dir + 'm-d.csv'\n",
    "        m_d = pd.read_csv(m_d_path, header=None).values\n",
    "        n_miRNA,n_disease = m_d.shape[0],m_d.shape[1]\n",
    "    else:\n",
    "        miRNA_sim, disease_sim = cal_faulty_sim(args)\n",
    "    \n",
    "        features = dict()\n",
    "        disease_sim_tensor = t.FloatTensor(disease_sim)\n",
    "        disease_sim_tensor = disease_sim_tensor.to(device)        \n",
    "        features['d'] = disease_sim_tensor#, 'edge_index': dd_edge_index}\n",
    "\n",
    "        miRNA_sim_tensor = t.FloatTensor(miRNA_sim)\n",
    "        miRNA_sim_tensor = miRNA_sim_tensor.to(device)\n",
    "        features['m'] = miRNA_sim_tensor#, 'edge_index': mm_edge_index}\n",
    "\n",
    "        features = train_features(features,args)\n",
    "        n_miRNA,n_disease = miRNA_sim.shape[0],disease_sim.shape[0]\n",
    "\n",
    "    seal_data = SEALData_case_study(g=graph, split_edge=split_edge, n_miRNA=n_miRNA, n_disease=n_disease, hop=args.hop,\\\n",
    "                         neg_samples=args.neg_samples, subsample_ratio=args.subsample_ratio, use_coalesce=False,\\\n",
    "                         random_seed=args.randseed, prefix=args.dataset+str(args.randseed), \\\n",
    "                         save_dir=args.save_dir, num_workers=args.num_workers, print_fn=print)\n",
    "    num_nodes = graph.num_nodes()\n",
    "    node_attribute = None #jicm\n",
    "    edge_weight = None #jicm\n",
    "\n",
    "    train_data = seal_data('train')\n",
    "    test_data = seal_data('test')#only use the #miRNA for disease-id\n",
    "    print(\"len of train_data: \", len(train_data), \",len of test data: \", len(test_data))\n",
    "\n",
    "    train_loader = GraphDataLoader(train_data, batch_size=args.batch_size)\n",
    "    test_loader = GraphDataLoader(test_data, batch_size=args.batch_size)\n",
    "    if  'irgat-gcn' in args.method:\n",
    "        model = GCN(num_layers=args.num_layers,\n",
    "                    hidden_units=args.hidden_units,\n",
    "                    gcn_type=args.gcn_type,\n",
    "                    pooling_type=args.pooling,\n",
    "                    node_attributes=features,#node_attribute,\n",
    "                    edge_weights=edge_weight,\n",
    "                    node_embedding=None,\n",
    "                    use_embedding=False,\n",
    "                    num_nodes=num_nodes,\n",
    "                    dropout=args.dropout)\n",
    "    elif  'irgat-dgcnn' in args.method:\n",
    "        model = DGCNN(num_layers=args.num_layers, \n",
    "                    hidden_units=args.hidden_units,\n",
    "                    k=args.sort_k,\n",
    "                    gcn_type=args.gcn_type,\n",
    "                    node_attributes=features,\n",
    "                    edge_weights=edge_weight,\n",
    "                    node_embedding=None,\n",
    "                    use_embedding=False,\n",
    "                    num_nodes=num_nodes,\n",
    "                    dropout=args.dropout)\n",
    "    else:\n",
    "        print('Invalid method name, please input the right name!')\n",
    "        return\n",
    "    \n",
    "    model = model.to(device)\n",
    "    parameters = model.parameters()\n",
    "    optimizer  = t.optim.Adam(parameters, lr=args.lr)#, weight_decay=args.decay)\n",
    "    scheduler  = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min',factor=0.1, patience=4, verbose=True) \n",
    "    \n",
    "    regression_crit = Myloss()\n",
    "    model = model.to(device)\n",
    "    regression_crit = regression_crit.to(device)\n",
    "    \n",
    "    model.train()\n",
    "    loss_fn = BCEWithLogitsLoss()\n",
    "    for epoch in range(0, args.epochs):    \n",
    "        total_loss = 0\n",
    "        pbar = tqdm(train_loader,ncols=100)\n",
    "        for i, (g, labels) in enumerate(pbar):\n",
    "            g = g.to(device)\n",
    "            labels = labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            logits = model(g, g.ndata['z'], g.ndata[NID], g.edata[EID])\n",
    "            loss = loss_fn(logits, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss += loss.item() * args.batch_size#\n",
    "            pbar.set_postfix({'loss:':loss.item()})\n",
    "        \n",
    "        scheduler.step(total_loss)\n",
    "   \n",
    "    test_pred_lbl, test_pair_lbl = evaluate(model=model,dataloader=test_loader,device=device)#jicm \n",
    "        \n",
    "    #save scores\n",
    "    if args.save_score:\n",
    "        score_save_dir = args.result_dir + str(args.randseed)\n",
    "        score_save_dir = standardize_dir(score_save_dir)\n",
    "        score_path = score_save_dir  + 'case_'+str(dis_id) + '_' + args.method + '_' + args.sim_type  +'.csv'\n",
    "        save_scores(test_pred_lbl.tolist(), test_pair_lbl.tolist(), score_path)   \n",
    "\n",
    "def query_topK(args,dis_id=236, k=50):\n",
    "    score_save_dir = args.result_dir + str(args.randseed)\n",
    "    score_save_dir = standardize_dir(score_save_dir)\n",
    "    score_path = score_save_dir  + 'case_'+str(dis_id) + '_' + args.method + '_' + args.sim_type  +'.csv'\n",
    "    score = pd.read_csv(score_path, header=None)\n",
    "\n",
    "    mirna_name =  pd.read_csv(osp.join(args.data_dir, 'mirna_name.csv'),header = None)[1].tolist()\n",
    "    disease    = pd.read_csv(osp.join(args.data_dir, 'disease_name.csv'),header = None)[1].tolist()\n",
    "    n_miRNA, n_disease = len(mirna_name), len(disease)\n",
    "    y_pred, y_true = score[0].tolist(), score[1].tolist()\n",
    "\n",
    "    loss_all = {}\n",
    "    loss = [] #for one disease, (loss, #mirnas)\n",
    "    for i in range(n_miRNA):#all mirna \n",
    "        loss_all[i] = y_pred[i]\n",
    "        if y_true[i] > 0.5:\n",
    "            loss.append(y_pred[i])\n",
    "\n",
    "    cc=0\n",
    "    loss_all = sorted(loss_all.items(), key=lambda d:d[1], reverse = True)\n",
    "    pred_K = [mirna_name[ls[0]]for ls in loss_all[:k]]\n",
    "    pred_K_path=args.result_dir+'case_dis_'+str(dis_id)+'_predK.csv'\n",
    "    save2File(pred_K, pred_K_path)\n",
    "\n",
    "    print(\"disease no: \", dis_id)\n",
    "    for i in range(k):\n",
    "        if loss_all[i][1] in loss:\n",
    "            cc = cc + 1\n",
    "            print(mirna_name[loss_all[i][0]])\n",
    "        else:\n",
    "            print(mirna_name[loss_all[i][0]], 'unknown')\n",
    "    print(\"test disease no:{},matched:{}/{}\".format(dis_id,cc,k))\n",
    "\n",
    "def verify_case(pred_K, bench_in_one_dis):\n",
    "    '''pred_K: for one disease, the top K predicted miRNAs\n",
    "       bench_in_one_dis: for one disease, the dbdmec/miR2disease verificated assoc\n",
    "    '''\n",
    "    pred_rna  = pd.read_csv(pred_K,header = None)[0].tolist()\n",
    "    bench_rna  = pd.read_csv(bench_in_one_dis,header = None)[0].tolist()\n",
    "    bench_rna = [str(rna).lower() for rna in bench_rna]\n",
    "\n",
    "    cc=0\n",
    "    for i in range(len(pred_rna)):\n",
    "        rna = pred_rna[i]\n",
    "        for brna in bench_rna:\n",
    "            if rna in brna:\n",
    "                print(\"confirmed\")\n",
    "                cc = cc + 1\n",
    "                break\n",
    "        else:\n",
    "            print(\"unknown\")   \n",
    "    print(\"{}/{} miRNA confirmed\".format(cc,len(pred_rna)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "hmdd = 'HMDD3.2/'\n",
    "\n",
    "parser = argparse.ArgumentParser(description='Subgraph neural networks learning for miRNA-disease association prediction')\n",
    "parser.add_argument('--epochs', type=int, default=130, metavar='N', help='number of epochs to train')\n",
    "parser.add_argument('--data_dir', default='../data/'+hmdd, help='dataset directory')\n",
    "parser.add_argument('--dataset', default=hmdd, help='dataset name')\n",
    "parser.add_argument('--fold_dir', default='../data/'+hmdd+'folds/', help='dataset directory')\n",
    "parser.add_argument('--result_dir', default='../data/'+hmdd+'results_subsample_1.0/', help='saved result directory')\n",
    "parser.add_argument('--method', default='irgat-dgcnn', help='method')\n",
    "parser.add_argument('--save_score', default=True, help='whether to save the predicted score or not')\n",
    "parser.add_argument('--sim_type', default='functional1', help='the miRNA and disease sim, pass in \"functional2\" for miRNA functional + disease semantic(with phenotype info added),'\n",
    "                                                              '\"none\" for none miRNA and disease additional info used,'\n",
    "                                                              '\"functional1\" for miRNA functional and disease semantic only,'\n",
    "                                                              '\"gip\" for miRNA and disease GIP kernel similarity,'\n",
    "                                                              '\"seq\" for miRNA sequence and disease semantic')\n",
    "parser.add_argument('--randseed', default=112, help='the random seed')\n",
    "\n",
    "parser.add_argument('--num_layers', type=int, default=3)\n",
    "parser.add_argument('--hidden_units', type=int, default=32)\n",
    "parser.add_argument('--gcn_type', type=str, default='gcn')\n",
    "parser.add_argument('--pooling', type=str, default='sum')\n",
    "parser.add_argument('--dropout', type=str, default=0.5)\n",
    "parser.add_argument('--lr', type=float, default=0.0001)\n",
    "parser.add_argument('--decay', type=float, default=0.0001)\n",
    "parser.add_argument('--hop', type=int, default=1)\n",
    "parser.add_argument('--sort_k', type=int, default=20)\n",
    "parser.add_argument('--neg_samples', type=int, default=1,help='negative samples per positive sample')\n",
    "parser.add_argument('--subsample_ratio', type=float, default=1.)\n",
    "parser.add_argument('--batch_size', type=int, default=64)\n",
    "parser.add_argument('--num_workers', type=int, default=8)\n",
    "parser.add_argument('--random_seed', type=int, default=123)\n",
    "parser.add_argument('--save_dir', type=str, default='./processed')\n",
    "args = parser.parse_args(args=[])\n",
    "\n",
    "   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dis_id =1 #1:colon neoplasms\n",
    "args.sim_type = 'functional2'\n",
    "case_study(dis_id,args)\n",
    "\n",
    "query_topK(args,dis_id=dis_id, k=50)\n",
    "pred_K, bench_in_one_dis = args.result_dir+'case_dis_'+str(dis_id)+'_predK.csv', args.result_dir+'case_dis_'+str(dis_id)+'_dbdemc.csv'\n",
    "verify_case(pred_K, bench_in_one_dis)\n",
    "bench_in_one_dis = args.result_dir+'case_dis_'+str(dis_id)+'_dbdemc.csv'\n",
    "verify_case(pred_K, bench_in_one_dis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jtorch-dgl",
   "language": "python",
   "name": "jtorch-dgl"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
